#!/usr/bin/python3

# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import os
import re
import struct
import parted
import subprocess
from parted import _ped
from .util import Util, PartiUtil
from .const import PartitionTableType, FileSystemType


class Handy:

    @staticmethod
    def getDiskDevPathList():
        ret = []
        for line in subprocess.check_output(["lsblk", "-o", "NAME,TYPE", "-n"], universal_newlines=True).split("\n"):
            m = re.fullmatch("(\\S+)\\s+(\\S+)", line)
            if m is None:
                continue
            if m.group(2) != "disk":
                continue
            devPath = "/dev/" + m.group(1)
            if not PartiUtil.inScope(devPath):
                continue
            if not os.access(devPath, os.R_OK):
                continue
            ret.append(devPath)
        return ret

    @staticmethod
    def getDataDiskParam(devPath):
        # get partition table type
        partiTableType = None
        try:
            partedDev = parted.getDevice(devPath)
            partedDisk = parted.newDisk(partedDev)
            partiTableType = PartitionTableType(partedDisk.type)
        except _ped.DiskException:
            return None
        except ValueError:
            return None

        # get filesystem information
        fs = None
        fsUuid = None
        label = None
        try:
            fsType, fsUuid, label = Util.getBlkDevFsTypeAndFsUuidAndLabel(PartiUtil.diskToParti(devPath, 1))
            fs = FileSystemType(fsType)
        except ValueError:
            return None

        # only 1 partition is allowed
        if PartiUtil.diskHasMoreParti(devPath, 1):
            return None

        # partition type specific check
        if partiTableType == PartitionTableType.GPT:
            if not Handy.checkGptDisk(devPath):
                return None
        elif partiTableType == PartitionTableType.MBR:
            if not Handy.checkMbrDisk(devPath):
                return None
        else:
            assert False

        return (partiTableType, fs, fsUuid, label)

    @staticmethod
    def checkGptDisk(devPath):
        # struct mbr_partition_record {
        #     uint8_t  boot_indicator;
        #     uint8_t  start_head;
        #     uint8_t  start_sector;
        #     uint8_t  start_track;
        #     uint8_t  os_type;
        #     uint8_t  end_head;
        #     uint8_t  end_sector;
        #     uint8_t  end_track;
        #     uint32_t starting_lba;
        #     uint32_t size_in_lba;
        # };
        mbrPartitionRecordFmt = "8BII"
        assert struct.calcsize(mbrPartitionRecordFmt) == 16

        # struct mbr_header {
        #     uint8_t                     boot_code[440];
        #     uint32_t                    unique_mbr_signature;
        #     uint16_t                    unknown;
        #     struct mbr_partition_record partition_record[4];
        #     uint16_t                    signature;
        # };
        mbrHeaderFmt = "440sIH%dsH" % (struct.calcsize(mbrPartitionRecordFmt) * 4)
        assert struct.calcsize(mbrHeaderFmt) == 512

        # struct gpt_entry {
        #     struct gpt_guid type;
        #     struct gpt_guid partition_guid;
        #     uint64_t        lba_start;
        #     uint64_t        lba_end;
        #     uint64_t        attrs;
        #     uint16_t        name[GPT_PART_NAME_LEN];
        # };
        gptEntryFmt = "16s16sQQQ36H"
        assert struct.calcsize(gptEntryFmt) == 128

        # struct gpt_header {
        #     uint64_t            signature;
        #     uint32_t            revision;
        #     uint32_t            size;
        #     uint32_t            crc32;
        #     uint32_t            reserved1;
        #     uint64_t            my_lba;
        #     uint64_t            alternative_lba;
        #     uint64_t            first_usable_lba;
        #     uint64_t            last_usable_lba;
        #     struct gpt_guid     disk_guid;
        #     uint64_t            partition_entry_lba;
        #     uint32_t            npartition_entries;
        #     uint32_t            sizeof_partition_entry;
        #     uint32_t            partition_entry_array_crc32;
        #     uint8_t             reserved2[512 - 92];
        # };
        gptHeaderFmt = "QIIIIQQQQ16sQIII420s"
        assert struct.calcsize(gptHeaderFmt) == 512

        dev = parted.getDevice(devPath)

        # FIXME: we can't use self._partedReadSectors() since it returns str, not bytes, what a bug!
        # mbrHeader = struct.unpack(mbrHeaderFmt, self._partedReadSectors(dev, 0, 1)[:struct.calcsize(mbrHeaderFmt)])

        with open(devPath, "rb") as f:
            # read Protective MBR header
            buf = f.read(struct.calcsize(mbrHeaderFmt))
            mbrHeader = struct.unpack(mbrHeaderFmt, buf)

            # check Protective MBR header
            if not Util.isBufferAllZero(mbrHeader[0]):
                return False
            if mbrHeader[1] != 0:
                return False
            if mbrHeader[2] != 0:
                return False
            if mbrHeader[4] != 0xAA55:
                return False

            # check Protective MBR Partition Record
            pRec = struct.unpack_from(mbrPartitionRecordFmt, mbrHeader[3], 0)
            if pRec[4] != 0xEE:
                return False
            if pRec[0] != 0:
                return False

            # other Partition Record should be filled with zero
            if not Util.isBufferAllZero(mbrHeader[struct.calcsize(mbrPartitionRecordFmt):]):
                return False

            # read to gpt header
            buf = f.read(dev.sectorSize - struct.calcsize(mbrHeaderFmt))
            if not Util.isBufferAllZero(buf):
                return False

            # ghnt and check primary and backup GPT header
            pass

            # get the first GPT partition GUID
            gptHeader = struct.unpack(gptHeaderFmt, f.read(struct.calcsize(gptHeaderFmt)))
            f.seek(gptHeader[10] * dev.sectorSize)
            partEntry = struct.unpack(gptEntryFmt, f.read(struct.calcsize(gptEntryFmt)))
            if partEntry[0] != Util.newGuid("EBD0A0A2-B9E5-4433-87C0-68B6B72699C7"):         # GUID of Microsoft Basic Data
                return False

        # all passed
        return True

    @staticmethod
    def checkMbrDisk(devPath):
        # FIXME
        # no boot code is allowed
        # with open(devPath, "rb") as f:
        #     if not Util.isBufferAllZero(f.read(440)):
        #         return False

        # all passed
        return True

    @staticmethod
    def _partedReadSectors(partedDev, startSector, sectorCount):
        partedDev.open()
        try:
            return partedDev.read(startSector, sectorCount)
        finally:
            partedDev.close()
